package org.example.lys.utils;

import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.ImageType;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.springframework.util.CollectionUtils;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author ; lidongdong
 * @description PDF转png 商户合同OCR识别
 * @date 2021-03-05
 */
public class PdfToImgUtil {
    /***
     * PDF文件转PNG图片，全部页数
     *s
     * @param pdfFilemultipartFile pdf文件
     * @param imgFilePath 图片完整路径
     * @return
     */
    public static File pdf2Image(MultipartFile pdfFilemultipartFile, String imgFilePath) throws IOException {
        return pdf2Image(pdfFilemultipartFile.getInputStream(), imgFilePath, 130, new ArrayList<>());
    }

    /***
     * PDF文件转PNG图片，全部页数
     *s
     * @param pdfFilemultipartFile pdf文件
     * @param imgFilePath 图片完整路径
     * @param assign 需要转换的页码，其他的将被丢弃
     * @return
     */
    public static File pdf2Image(MultipartFile pdfFilemultipartFile, String imgFilePath, List<Integer> assign) throws IOException {
        return pdf2Image(pdfFilemultipartFile.getInputStream(), imgFilePath, 130, assign);
    }

    /***
     * PDF文件转PNG图片，全部页数
     *
     * @param pdfFilePath pdf完整路径
     * @param imgFilePath 图片完整路径
     * @return
     */
    public static File pdf2Image(String pdfFilePath, String imgFilePath) throws IOException {
        return pdf2Image(pdfFilePath, imgFilePath, 130, new ArrayList<>());
    }

    /***
     * PDF文件转PNG图片，全部页数
     *s
     * @param pdfFilePath pdf完整路径
     * @param imgFilePath 图片完整路径
     * @param dpi dpi越大转换后越清晰，相对转换速度越慢
     * @return
     */
    public static File pdf2Image(String pdfFilePath, String imgFilePath, int dpi) throws IOException {
        return pdf2Image(pdfFilePath, imgFilePath, dpi, new ArrayList<>());
    }

    /***
     * PDF文件转PNG图片
     *s
     * @param pdfFilePath pdf完整路径
     * @param imgFilePath 图片完整路径
     * @param dpi dpi越大转换后越清晰，相对转换速度越慢
     * @param assign 需要转换的页码，其他的将被丢弃
     * @return
     */
    public static File pdf2Image(String pdfFilePath, String imgFilePath, int dpi, List<Integer> assign) throws IOException {
        File file = new File(pdfFilePath);
        return pdf2Image(new FileInputStream(file), imgFilePath, dpi, assign);
    }

    /***
     * PDF文件转PNG图片，全部页数
     *s
     * @param pdfFileStream pdf文件流
     * @param imgFilePath 图片完整路径
     * @param dpi dpi越大转换后越清晰，相对转换速度越慢
     * @param assign 需要转换的页码，其他的将被丢弃
     * @return
     */
    public static File pdf2Image(InputStream pdfFileStream, String imgFilePath, int dpi, List<Integer> assign) throws IOException {

        PDDocument pdDocument;
        pdDocument = PDDocument.load(pdfFileStream);
        PDFRenderer renderer = new PDFRenderer(pdDocument);
        /* dpi越大转换后越清晰，相对转换速度越慢 */
        List<BufferedImage> piclist = new ArrayList<>();
        for (int i = 0; i < pdDocument.getNumberOfPages(); i++) {
            if (!CollectionUtils.isEmpty(assign) && !assign.contains(i)) {
                continue;
            }
            BufferedImage image = renderer.renderImageWithDPI(i, dpi);
            image.flush();
            piclist.add(image);
        }
        return yPic(piclist, imgFilePath);
    }

    /**
     * 将宽度相同的图片，竖向追加在一起 ##注意：宽度必须相同
     *
     * @param piclist 文件流数组
     * @param outPath 输出路径
     */
    public static File yPic(List<BufferedImage> piclist, String outPath) throws IOException {// 纵向处理图片
        if (piclist == null || piclist.size() <= 0) {
            System.out.println("图片数组为空!");
            return null;
        }
        int height = 0, // 总高度
                width = 0, // 总宽度
                _height = 0, // 临时的高度 , 或保存偏移高度
                __height = 0, // 临时的高度，主要保存每个高度
                picNum = piclist.size();// 图片的数量
        int[] heightArray = new int[picNum]; // 保存每个文件的高度
        BufferedImage buffer = null; // 保存图片流
        List<int[]> imgRGB = new ArrayList<int[]>(); // 保存所有的图片的RGB
        int[] _imgRGB; // 保存一张图片中的RGB数据
        for (int i = 0; i < picNum; i++) {
            buffer = piclist.get(i);
            heightArray[i] = _height = buffer.getHeight();// 图片高度
            if (i == 0) {
                width = buffer.getWidth();// 图片宽度
            }
            height += _height; // 获取总高度
            _imgRGB = new int[width * _height];// 从图片中读取RGB
            _imgRGB = buffer.getRGB(0, 0, width, _height, _imgRGB, 0, width);
            imgRGB.add(_imgRGB);
        }
        _height = 0; // 设置偏移高度为0
        // 生成新图片
        BufferedImage imageResult = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
        for (int i = 0; i < picNum; i++) {
            __height = heightArray[i];
            if (i != 0) {
                _height += __height; // 计算偏移高度
            }
            imageResult.setRGB(0, _height, width, __height, imgRGB.get(i), 0, width); // 写入流中
        }
        File outFile = new File(outPath);
        ImageIO.write(imageResult, "png", outFile);// 写图片
        return outFile;
    }

    public static Map<Integer, File> pdfPathToImagePaths(String pdfFilePath, String outPath) throws IOException {
        return pdfPathToImagePaths(new FileInputStream(new File(pdfFilePath)), outPath, new ArrayList<>());
    }


    public static Map<Integer, File> pdfPathToImagePaths(InputStream pdfInputStream, String outPath) throws IOException {
        return pdfPathToImagePaths(pdfInputStream, outPath, new ArrayList<>());
    }

    /**
     * 将多页pdf转化为多张图片
     *
     * @param pdfInputStream 表示pdf的路径
     * @param assign         需要转换的页码，其他的将被丢弃
     * @return 转化后的图片的路径集合
     * @throws IOException
     */
    public static Map<Integer, File> pdfPathToImagePaths(InputStream pdfInputStream, String outPath, List<Integer> assign) throws IOException {
        PDDocument pdDocument = PDDocument.load(pdfInputStream);
        int pageCount = pdDocument.getNumberOfPages();
        PDFRenderer pdfRenderer = new PDFRenderer(pdDocument);
        Map<Integer, File> imagePathList = new HashMap<>();
        for (int pageIndex = 0; pageIndex < pageCount; pageIndex++) {
            if (!CollectionUtils.isEmpty(assign) && !assign.contains(pageIndex)) {
                continue;
            }
            String imgPath = outPath + File.separator + pageIndex + ".png";
            /* dpi越大转换后越清晰，相对转换速度越慢 */
            BufferedImage image = pdfRenderer.renderImageWithDPI(pageIndex, 100, ImageType.RGB);
            File file = new File(imgPath);
            ImageIO.write(image, "png", file);
            imagePathList.put(pageIndex, file);
        }
        pdDocument.close();
        return imagePathList;
    }

//    public static void main(String[] args) throws IOException {
//        String filePath = "/Users/lidongdong/Desktop/商户签约合同210305.pdf";
////        pdf2Image(filePath, "/Users/lidongdong/Desktop/商户签约合同210308.png", 80, Arrays.asList(3, 6, 8));
//        System.out.println(pdfPathToImagePaths(new FileInputStream(new File(filePath)), "/Users/lidongdong/Desktop/",
//                Arrays.asList(1, 2, 6, 8)).get(1).getPath());
//    }

}
